import sys
import numpy as np 
import os
import time
import math
from PIL import Image
import cv2
from datetime import datetime
from pynq import Xlnk
from pynq import Overlay
import pynq
import struct
from multiprocessing import Process, Pipe, Queue, Event, Manager
from IoU import Average_IoU
import cv2 as cv
from pynq.lib.video import *

anchor = [1.4940052559648322, 2.3598481287086823, 4.0113013115312155, 5.760873975661669]
bbox_m = [52., 48., 28., 30., 124., 47., 52., 23., 23., 125.]
qm = 131072.0
w = 40
h = 20
# IMG_DIR = '../sample1000/'
IMG_DIR_2 = '../my_data/dataset_training/riding4/'
IMG_DIR_0 = '../my_data/dataset_training/boat5/'
IMG_DIR_3 = '../my_data/dataset_training/car11/'
IMG_DIR_1 = '../my_data/dataset_training/truck1/'
file_type = '.jpg'
READ_IMAGE_START = 0;
MAX_FRAME = 900
pixel_bits = 24
DP_width = 1920
DP_height = 1080
photo_height = 360
photo_width = 640
p1_x_min = math.floor((DP_width-photo_width*2)/2)
p1_x_max = p1_x_min + photo_width
p1_y_min = math.floor((DP_height-photo_height*2)/2)
p1_y_max = p1_y_min + photo_height

p2_x_min = p1_x_max
p2_x_max = p2_x_min + photo_width
p2_y_min = p1_y_min
p2_y_max = p1_y_max

p3_x_min = p1_x_min
p3_x_max = p1_x_max
p3_y_min = p1_y_max
p3_y_max = p3_y_min + photo_height

p4_x_min = p1_x_max
p4_x_max = p4_x_min + photo_width
p4_y_min = p1_y_max
p4_y_max = p4_y_min + photo_height

DP_X_MIN = [p1_x_min, p2_x_min, p3_x_min, p4_x_min]
DP_X_MAX = [p1_x_max, p2_x_max, p3_x_max, p4_x_max]
DP_Y_MIN = [p1_y_min, p2_y_min, p3_y_min, p4_y_min]
DP_Y_MAX = [p1_y_max, p2_y_max, p3_y_max, p4_y_max]

BATCH_SIZE = 4

def sigmoid(x):
    return 1/(1+np.exp(-x))

def stitch(idx):
    
    show_image   = np.zeros((4,360,640,3),np.uint8)

    start = time.time()

    frame0 = Image.open(IMG_DIR_0+str(idx).zfill(4) + file_type)
    frame1 = Image.open(IMG_DIR_1+str(idx).zfill(4) + file_type)
    frame2 = Image.open(IMG_DIR_2+str(idx).zfill(4) + file_type)
    frame3 = Image.open(IMG_DIR_3+str(idx).zfill(4) + file_type)

    end = time.time()
    read_time = (end - start)

    frame_count_add = 1

    image           = np.zeros((160,320,16),np.uint8)
    image[:,:,0:3]  = np.array(frame0.resize((320, 160)))
    image[:,:,3:6]  = np.array(frame1.resize((320, 160)))
    image[:,:,6:9]  = np.array(frame2.resize((320, 160)))
    image[:,:,9:12] = np.array(frame3.resize((320, 160)))
    
    end2 = time.time()
    
    show_image[0] = np.array(frame0)
    show_image[1] = np.array(frame1)
    show_image[2] = np.array(frame2)
    show_image[3] = np.array(frame3)
    
    read_time2 = (end2 - end)

    return image, show_image, read_time, frame_count_add, read_time2

def compute_bounding_box(bbox, bbox_origin, idx, result):
    show_image_bbox = np.zeros((4,4),int)
    for b in range(4):
        if(bbox_origin[b,4]>0):
            xs = bbox_origin[b][0]*bbox_m[5]/qm
            ys = bbox_origin[b][1]*bbox_m[6]/qm
            ws = bbox_origin[b][2]*bbox_m[7]/qm
            hs = bbox_origin[b][3]*bbox_m[8]/qm
            ws_inb = math.exp(ws)*anchor[2]
            hs_inb = math.exp(hs)*anchor[3]
        else:
            xs = bbox_origin[b][0]*bbox_m[0]/qm
            ys = bbox_origin[b][1]*bbox_m[1]/qm
            ws = bbox_origin[b][2]*bbox_m[2]/qm
            hs = bbox_origin[b][3]*bbox_m[3]/qm
            ws_inb = np.exp(ws)*anchor[0]
            hs_inb = np.exp(hs)*anchor[1]
        xs_inb = sigmoid(xs) + bbox_origin[b][5]
        ys_inb = sigmoid(ys) + bbox_origin[b][6]
        bcx = xs_inb/w
        bcy = ys_inb/h
        bw = ws_inb/w
        bh = hs_inb/h
        bbox[b][0] = bcx - bw/2.0
        bbox[b][1] = bcy - bh/2.0
        bbox[b][2] = bbox[b][0] + bw
        bbox[b][3] = bbox[b][1] + bh

        x1 = int(round(bbox[b][0] * 640))
        y1 = int(round(bbox[b][1] * 360))
        x2 = int(round(bbox[b][2] * 640))
        y2 = int(round(bbox[b][3] * 360))
        x1 = np.clip(x1,1,640)
        y1 = np.clip(y1,1,360)
        x2 = np.clip(x2,1,640)
        y2 = np.clip(y2,1,360)
        
        result.write(str(idx).zfill(4) + file_type +' '+str([x1, x2, y1, y2])+'\n')
        #print(batch[b], [x1, x2, y1, y2])  
        show_image_bbox[b] = [x1, x2, y1, y2]
    return show_image_bbox

def display_screen(show_image_buff, show_image_bbox):
    frame = displayport.newframe()
    for i in range(4):
        x1 = show_image_bbox[i][0]
        x2 = show_image_bbox[i][1]
        y1 = show_image_bbox[i][2]
        y2 = show_image_bbox[i][3]
        cv2.rectangle(show_image_buff[i], (x1,y1), (x2,y2), (0,255,0), 4)
        frame[DP_Y_MIN[i]:DP_Y_MAX[i],DP_X_MIN[i]:DP_X_MAX[i],:] = show_image_buff[i]
    displayport.writeframe(frame)
    #videoWriter.write(frame)

#4

################################## Init FPGA ##################################
xlnk = Xlnk()
xlnk.xlnk_reset()

img    = xlnk.cma_array(shape=[160,320,16], dtype=np.uint8)
fm     = xlnk.cma_array(shape=(628115*32), dtype=np.uint8)
weight = xlnk.cma_array(shape=(220672),  dtype=np.int16)
biasm  = xlnk.cma_array(shape=(432*16),  dtype=np.int16)

print("Allocating memory done")

parameter = np.fromfile("SkyNet.bin", dtype=np.int16)
np.copyto(weight, parameter[0:220672])
np.copyto(biasm[0:428*16], parameter[220672:])
print("Parameters loading done")

overlay = Overlay("SkyNet.bit")
print("Bitstream loaded")
if len(sys.argv)>1:
    os.system("sudo ./dfs 0 "+sys.argv[1])

SkyNet = overlay.SkyNet
SkyNet.write(0x10, img.physical_address)
SkyNet.write(0x1c, fm.physical_address)
SkyNet.write(0x28, weight.physical_address)
SkyNet.write(0x34, biasm.physical_address)

rails    = pynq.get_rails()
recorder = pynq.DataRecorder(rails["power1"].power)

#5
################################## Main process ##################################
bbox_origin = np.empty(64, dtype=np.int16)
bbox        = np.zeros((4,4),dtype=np.float32)
result      = open('predict.txt','w+')
# prepare display
displayport = DisplayPort()
displayport.configure(VideoMode(DP_width, DP_height, pixel_bits), PIXEL_RGB)

fps = 20
size=(1920,1080)

image_buff       = np.zeros((160,320,16),np.uint8)
show_image_buff  = np.zeros((4,360,640,3),np.uint8)
total_read_time = 0
total_stitch_time = 0

print("\n**** Start to detect")
all_start = time.time()
frame_count = 0
idx = 0
with recorder.record(0.05):
    #for idx in range(1,MAX_FRAME):
    while(True):
        idx = idx + 1
        if(idx > 900):
            idx = 1
        if (idx == 0):
            image_buff,show_image_buff,read_time,frame_count_add,read_time2  = stitch(idx)
            np.copyto(img, image_buff)
            total_read_time += read_time
            frame_count += frame_count_add
            total_stitch_time += read_time2
        else:
            np.copyto(img, image_buff) 

        SkyNet.write(0x00, 1)

        if (0 <= idx < MAX_FRAME - 1):
            # pass
            image_buff,show_image_buff,read_time,frame_count_add, read_time2 = stitch(idx+1)
            total_read_time += read_time
            frame_count += frame_count_add
            total_stitch_time += read_time2
        if (0 < idx <= MAX_FRAME - 1):
            show_image_bbox = compute_bounding_box(bbox ,bbox_origin.reshape(4,16), idx, result)
            display_screen(show_image_buff, show_image_bbox)

        isready = SkyNet.read(0x00)
        while( isready == 1 ):
            isready = SkyNet.read(0x00)

        np.copyto(bbox_origin, biasm[428*16:])

        if (idx == MAX_FRAME - 1):
            show_image_bbox = compute_bounding_box(bbox ,bbox_origin.reshape(4,16), idx, result)
            display_screen(show_image_buff, show_image_bbox)

result.close()

all_end = time.time()
total_time = all_end - all_start
print("**** Detection finished\n")
print('Total time: ' + str(total_time) + ' s')
print('Total read time: ' + str(total_read_time) + ' s')
print('Total stitch time: ' + str(total_stitch_time) + ' s')
print('frame_count: ' + str(frame_count))
